package com.smart.milvus.bean;

import com.smart.milvus.conditions.query.BaseQueryWrapper;
import com.smart.milvus.conditions.search.BaseSearchWrapper;
import com.smart.milvus.exception.MilvusException;
import com.smart.milvus.model.Page;
import io.milvus.client.MilvusClient;
import io.milvus.grpc.LoadState;

import java.io.Serializable;
import java.util.List;

public class MilvusMapperImpl<T extends MilvusModel<?>> implements BaseMilvusMapper<T>, Serializable {

    protected final Class<T> entityType;
    private final MilvusClientService milvusClientService;

    public MilvusMapperImpl(Class<T> entityType,
                            MilvusClientService milvusClientService) {
        this.entityType = entityType;
        this.milvusClientService = milvusClientService;
    }

    @Override
    public MilvusClientService getMilvusClientService() {
        return milvusClientService;
    }

    @Override
    public MilvusClient getMilvusClient() {
        return milvusClientService.getMilvusClient();
    }

    @Override
    public boolean insert(T vectorModel) throws MilvusException {
        return milvusClientService.insert(vectorModel);
    }

    @Override
    public boolean remove(Object pk) throws MilvusException {
        return milvusClientService.remove(pk, entityType);
    }

    @Override
    public boolean batchRemove(BaseQueryWrapper<T, ? extends BaseQueryWrapper<T, ?>> wrapper) throws MilvusException {
        wrapper.setEntityClass(entityType);
        return milvusClientService.batchRemove(wrapper);
    }

    @Override
    public void loadCollection() throws MilvusException {
        milvusClientService.loadCollection(entityType);
    }

    @Override
    public LoadState getLoadState() throws MilvusException {
        return milvusClientService.getLoadState(entityType);
    }

    @Override
    public Long getLoadProgress() throws MilvusException {
        return milvusClientService.getLoadProgress(entityType);
    }

    @Override
    public void releaseCollection() throws MilvusException {
        milvusClientService.releaseCollection(entityType);
    }

    @Override
    public boolean hasCollection() throws MilvusException {
        return milvusClientService.hasCollection(entityType);
    }

    @Override
    public List<T> search(BaseSearchWrapper<T, ?, ? extends BaseSearchWrapper<T, ?, ?>> wrapper) throws MilvusException {
        wrapper.setEntityClass(entityType);
        return milvusClientService.search(wrapper);
    }

    @Override
    public List<T> query(BaseQueryWrapper<T, ? extends BaseQueryWrapper<T, ?>> wrapper) throws MilvusException {
        wrapper.setEntityClass(entityType);
        return milvusClientService.query(wrapper);
    }

    @Override
    public Page<T> queryPage(BaseQueryWrapper<T, ? extends BaseQueryWrapper<T, ?>> wrapper, Long page, Long pageSize) throws MilvusException {
        wrapper.setEntityClass(entityType);
        return milvusClientService.queryPage(wrapper, page, pageSize);
    }
}
